import time
from typing import Union, Callable

from scripts.Log import Log
from scripts.actions.Context import Context
from scripts.config.Config import Config
from scripts.device.AdbDevice import AdbDevice
from scripts.actions.Action import Action


class CleanupAction(Action):
    KEY = Action.register('cleanup')
    DISPLAY_NAME = 'Cleanup'

    REBOOT_ONLY = Config.STEP_CLEANUP_REBOOT_ONLY

    def __init__(self, context: Context):
        super().__init__(CleanupAction.KEY, CleanupAction.DISPLAY_NAME, context, False)

    def _on_backup(self, progress: Union[None, Callable[[int, int, str], None]]):
        if progress:
            progress(0, 1, 'Finishing...')

        if CleanupAction.REBOOT_ONLY:
            Log.d('directly reboot due to reboot only')
            self.wait_for_device().reboot()
        else:
            # do unroot
            self._on_unroot()

            # disable developer options
            self._on_disable_developer_options()

            # disable adb
            self._on_disable_adb()

        if progress:
            progress(1, 1, 'Done [%s]' % self.duration_time_str())

    def _on_restore(self, progress: Union[None, Callable[[int, int, str], None]]):
        if progress:
            progress(0, 1, 'Finishing...')

        if CleanupAction.REBOOT_ONLY:
            Log.d('directly reboot due to reboot only')
            self.wait_for_device().reboot()
        else:
            # do unroot
            self._on_unroot()

            # disable developer options
            self._on_disable_developer_options()

            # disable adb
            self._on_disable_adb()

        if progress:
            progress(1, 1, 'Done [%s]' % self.duration_time_str())

    def _on_unroot(self):
        device = self.wait_for_device()
        with device.shell() as shell:
            # remove temp dir on device
            Log.d('remove temp dir')
            shell.rm(self.device_temp_dir())

            # enable-verity before unroot
            # device.enable_verity()

            # unroot
            ver = self.android_version()
            if ver == 9:
                Log.d('unroot9')
                self.unroot9(device)
            elif ver == 11:
                Log.d('unroot11')
                self.unroot11(device)
            else:
                Log.e('unable to unroot device with android version: %d' % ver)

    def _on_disable_developer_options(self):
        Log.d('wait for device to disable developer options')
        device = self.wait_for_device()
        with device.shell() as shell:
            shell.wait_for_boot_complete()

            # close developer options
            Log.d('close developer options')
            shell.settings_put('global', 'development_settings_enabled', '0')
            shell.settings_put('global', 'developer_options_show', '0')
            # shell.am_broadcast('com.android.settingslib.development.DevelopmentSettingsEnabler.SETTINGS_CHANGED')

    def _on_disable_adb(self):
        Log.d('wait for device to disable adb')
        time.sleep(5)
        device = self.wait_for_device()
        with device.shell() as shell:
            shell.wait_for_boot_complete()

            # disable adb
            Log.d('disable adb')
            shell.disable_adb()
            # device.am_broadcast('com.android.settingslib.development.AbstractEnableAdbController.ENABLE_ADB_STATE_CHANGED')

            # stop adb process
            Log.d('stop adb')
            AdbDevice.kill_server()

    def unroot9(self, adb_device: AdbDevice) -> bool:
        # reboot to fastboot
        adb_device.reboot_bootloader()

        # get fastboot device
        fastboot_device = self.wait_for_fastboot_device()

        # do device root
        result = fastboot_device.disable_root_9()

        # reboot
        fastboot_device.reboot()

        return result

    def unroot11(self, adb_device: AdbDevice) -> bool:
        # reboot to fastboot
        adb_device.reboot_bootloader()

        # get fastboot device
        fastboot_device = self.wait_for_fastboot_device()

        # do device root
        result = fastboot_device.disable_root_11()

        # reboot
        fastboot_device.reboot()

        return result
